import cv2
import numpy as np
import time, os, sys, math
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--img', type=str, default='./1328301915.jpg')
parser.add_argument('--cam_dist', type=float, default=300)

args = parser.parse_args()

coin_diameter_mm = 25

if not os.path.exists(args.img):
    print('input image not exists.')
    exit()

# class Finger(object):
#     def __init__(self, name, direction, tip, length):
#         self.name = name
#         self.direction = direction
#         self.tip = tip
#         self.length = length

def normalize(vec):
    return vec / np.linalg.norm(vec)

font = cv2.FONT_HERSHEY_SIMPLEX

img = cv2.imread(args.img)
# img = cv2.resize(img, (img.shape[1] // 2, img.shape[0] // 2))
# img = cv2.flip(img, 1)
#Blur the image
blur = cv2.GaussianBlur(img,(5,5), 0)

gray = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
circles = cv2.HoughCircles(gray, cv2.HOUGH_GRADIENT, dp=2.2, minDist=100,
                           param1=200, param2=100)

if circles is not None:
    # append radius to list of diameters (we don't bother to multiply by 2)
    for (x, y, r) in circles[0, :]:
        if r < img.shape[0] / 8:
            # cv2.circle(img, (x, y), r, (0,0,255), 2)
            coin_radius = r
            r += 5
            coin_roi = gray[int(y-r):int(y+r), int(x-r):int(x+r)]
            coin_canny = cv2.Canny(coin_roi, 100, 200)
            coin_contours, coin_hierarchy = cv2.findContours(coin_canny,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)  
            # coin_roi = coin_roi // 2
            # cv2.drawContours(coin_roi, coin_contours, 0, 255)
            cv2.drawContours(img, coin_contours, 0, (0,0,0), 2, offset=(int(x-r), int(y-r)))
            raw_img = img.copy()
            #fix camera view angle
            src_pts = coin_contours[0]
            src_pts = src_pts.reshape((src_pts.shape[0], 2))[::10, :].copy()
            dst_pts = src_pts.copy()
            center = np.array((x, y))
            for i in range(src_pts.shape[0]):
                src_pts[i] = src_pts[i] + np.array((x-r, y-r))
                r_vec = src_pts[i] - center
                dst_pts[i] = center + normalize(r_vec) * coin_radius
                

            homography, status = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, ransacReprojThreshold=2)

            # img = cv2.warpPerspective(img, homography, (img.shape[1], img.shape[0]))
            # blur = cv2.GaussianBlur(img,(5,5), 0)
            # gray = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
            cv2.circle(img, (int(x), int(y)), int(coin_radius), (0,0,255), 2)
            break

# cv2.imshow('raw', raw_img)
# cv2.imshow('homo', img)
# cv2.waitKey()
print('coin radius: %spx' % coin_radius)

#Convert to HSV color space
hsv = cv2.cvtColor(blur,cv2.COLOR_BGR2HSV)

#Create a binary image with where white will be skin colors and rest is black
lower_skin = np.array([0,20,50], dtype=np.uint8)
upper_skin = np.array([20,255,255], dtype=np.uint8)
mask2 = cv2.inRange(hsv, lower_skin, upper_skin)

#Kernel matrices for morphological transformation    
kernel_square = np.ones((11,11),np.uint8)
kernel_ellipse= cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))

#Perform morphological transformations to filter out the background noise
#Dilation increase skin color area
#Erosion increase skin color area
dilation = cv2.dilate(mask2,kernel_ellipse,iterations = 1)
erosion = cv2.erode(dilation,kernel_square,iterations = 1)    
dilation2 = cv2.dilate(erosion,kernel_ellipse,iterations = 1)    
filtered = cv2.medianBlur(dilation2,5)
kernel_ellipse= cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(8,8))
dilation2 = cv2.dilate(filtered,kernel_ellipse,iterations = 1)
kernel_ellipse= cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(5,5))
dilation3 = cv2.dilate(filtered,kernel_ellipse,iterations = 1)
median = cv2.medianBlur(dilation2,5)
ret,thresh = cv2.threshold(median,127,255,0)

# cv2.imshow('thresh', thresh)
#Find contours of the filtered frame
contours, hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)   

#Find Max contour area (Assume that hand is in the frame)
max_area=100
ci=0	
for i in range(len(contours)):
    cnt=contours[i]
    area = cv2.contourArea(cnt)
    if(area>max_area):
        max_area=area
        ci=i  
        
#Largest area contour 			  
cnts = contours[ci]

#Draw Contours
cv2.drawContours(img, [cnts], 0, (122,122, 0), 2)

boundx,boundy,boundw,boundh = cv2.boundingRect(cnts)
margin = 100
boundx = np.clip(boundx - margin, 0, img.shape[1])
boundw = np.clip(boundw + margin * 2, 0, img.shape[1])
boundy = np.clip(boundy - margin, 0, img.shape[0])
boundh = np.clip(boundh + margin * 2, 0, img.shape[0])

epsilon = 0.001 * cv2.arcLength(cnts, True)
cnts = cv2.approxPolyDP(cnts, epsilon, True)

cv2.drawContours(img, [cnts], 0, (123,123, 255), 2)
# cv2.imshow('raw', img)
# cv2.waitKey()

#Find moments of the largest contour
moments = cv2.moments(cnts)

#Central mass of first order moments
if moments['m00']!=0:
    cx = int(moments['m10']/moments['m00']) # cx = M10/M00
    cy = int(moments['m01']/moments['m00']) # cy = M01/M00
centerMass=(cx,cy)    

#Draw center mass
cv2.circle(img,centerMass,7,[100,0,255],2)
cv2.putText(img,'Center',tuple(centerMass),font,1,(255,255,255),2)

#Find convex hull
hull = cv2.convexHull(cnts)

#Find convex defects
hull2 = cv2.convexHull(cnts,returnPoints = False)
defects = cv2.convexityDefects(cnts,hull2)

if defects is not None:
    #Get defect points and draw them in the original image
    valid_defects = []
    valid_defects_index = []
    avg_finger_length = 0
    for i in range(defects.shape[0]):
        s,e,f,d = defects[i,0]
        start = cnts[s][0]
        end = cnts[e][0]
        far = cnts[f][0]

        # find length of all sides of triangle
        a = np.linalg.norm(end - start)
        b = np.linalg.norm(far - start)
        c = np.linalg.norm(end - far)
        s = (a+b+c)/2
        ar = math.sqrt(s*(s-a)*(s-b)*(s-c))

        #distance between point and convex hull
        d=(2*ar)/a

        # apply cosine rule here
        angle = math.acos((b**2 + c**2 - a**2)/(2*b*c)) * 57


        # ignore angles > 90 and ignore points very close to convex hull(they generally come due to noise)
        if angle <= 90 and d>30:
            valid_defects.append([start, end, far])
            valid_defects_index.append(i)
            avg_finger_length += np.linalg.norm(end - far)
            
    avg_finger_length /= len(valid_defects)
    for i in range(len(valid_defects)):
        if np.linalg.norm(valid_defects[i][0] - valid_defects[i - 1][1]) < avg_finger_length / 3:
            valid_defects[i][0] = np.int32((valid_defects[i][0] + valid_defects[i - 1][1]) / 2)
            valid_defects[i - 1][1] = valid_defects[i][0]
    
    defect_dist = np.zeros(len(valid_defects))
    for i in range(len(valid_defects)):
        defect_dist[i] = np.linalg.norm(valid_defects[i][2] - valid_defects[i-1][2])

    finger_thumb = defect_dist.argmax()
    defect_dist[finger_thumb] = 0
    left_hand = finger_thumb  + 1 == defect_dist.argmax()
    if left_hand:        
        step = 1
    else:
        finger_thumb -= 1
        step = -1

    # finger_thumb = (finger_thumb - step) % len(valid_defects)
    finger_name =['thumb', 'index', 'middle', 'ring', 'little']
    finger_poly = []
    for i in range(len(valid_defects)):
        ii = (finger_thumb + i * step) % len(valid_defects)
        start = tuple(valid_defects[ii][0])
        end = tuple(valid_defects[ii][1])
        far = tuple(valid_defects[ii][2])
        
        # cv2.line(img,start,far,[0,255,0],2)
        # cv2.line(img,far,end,[0,255,0],2)
        # cv2.line(img,start,end,[0,0, 255],2)
        # cv2.line(img, start, tuple(centerMass), [0, 0, 0], 2)
        # cv2.line(img, end, tuple(centerMass), [0, 0, 0], 2)
        # cv2.circle(img,start,5,[0,0,255],-1)
        # cv2.circle(img,end,5,[255,0,0],-1)
        # cv2.circle(img,far,5,[60 * i, 60 * i, 60 * i],-1)
        if left_hand:
            # if i == 3:
            #     cv2.putText(img, finger_name[-1], (end[0] + 10, end[1]), font,1,(255,255,255),2)

            # cv2.putText(img, finger_name[i], (start[0] + 10, start[1]), font,1,(255,255,255),2)
            finger_poly.append(start)
        else:
            # if i == 0:
            #     cv2.putText(img, finger_name[0], (end[0] + 10, end[1]), font,1,(255,255,255),2)

            # cv2.putText(img, finger_name[i + 1], (start[0] + 10, start[1]), font,1,(255,255,255),2)
            finger_poly.append(end)
        finger_poly.append(far)

    if left_hand:
        finger_poly.append(end)
        prev_defect = defects[valid_defects_index[finger_thumb] - 1, 0]
        next_defect = defects[(valid_defects_index[finger_thumb-1] + 1) % defects.shape[0], 0]
    else:
        finger_poly.append(start)
        prev_defect = defects[(valid_defects_index[finger_thumb] +1) % defects.shape[0], 0]
        next_defect = defects[(valid_defects_index[finger_thumb + 1] - 1), 0]

    finger_poly.insert(0, cnts[prev_defect[2]][0])
    finger_poly.append(cnts[next_defect[2]][0])
    finger_poly = np.array(finger_poly).reshape((1, len(finger_poly), 2))
    cv2.polylines(img, finger_poly, False, [0,255,0], 2)

    finger_directions = []
    finger_length = []
    for i in range(0, finger_poly.shape[1] - 2, 2):
        fstart = finger_poly[0, i, :] - finger_poly[0, i + 1, :]
        fend = finger_poly[0, i + 2, :] - finger_poly[0, i + 1, :]
        fsd = np.linalg.norm(fstart)
        fed = np.linalg.norm(fend)
        finger_directions.append( normalize(fstart / fsd + fend / fed))
        finger_length.append((fsd + fed) / 2)
        if i == 0:
            finger_length[-1] = (finger_length[-1] + avg_finger_length) / 2
        # fmiddle = finger_poly[0, i + 1, :] + np.int32(finger_directions[-1] * finger_length[-1])
        # cv2.line(img, tuple(finger_poly[0, i + 1, :]), tuple(fmiddle), (0,0,0), 2)

    num_step = 40
    #finetune finger directions
    dist, labels = cv2.distanceTransformWithLabels(thresh, cv2.DIST_L2, 5)
    dist = np.uint8(dist * .5)
    finger_test = np.zeros(dist.shape, dtype=np.uint8)
    intersect_result = dist.copy()
    cv2.drawContours(intersect_result, [cnts], 0, 255, 2)
    
    max_finger_half_width = avg_finger_length / 4
    finger_lines = []
    for i in range(5):
        step = finger_length[i] / (num_step + 1)
        finger_tip = finger_poly[0, 2*i + 1, :]
        finger_dir = finger_directions[i]
        perp_dir = finger_dir.copy()
        perp_dir[0], perp_dir[1] = perp_dir[1], -perp_dir[0]
        assert np.dot(perp_dir, finger_dir) == 0
        finger_middle_line = []
        for j in range(num_step + 2):
            finger_test[:,:] = 0        
            pos = finger_tip + finger_dir * ((j + 1) * step)
            cv2.line(finger_test, tuple(np.int32(pos + perp_dir * max_finger_half_width)),\
                tuple(np.int32(pos - perp_dir * max_finger_half_width)), 1)

            intersect = finger_test * dist
            middle_pt = np.array(np.unravel_index(intersect.argmax(), intersect.shape))
            middle_pt[0], middle_pt[1] = middle_pt[1], middle_pt[0]
            cv2.circle(intersect_result, tuple(middle_pt), 5, 255, 2)
            cv2.circle(img, tuple(middle_pt), 5, (255,255,255), 2)
            finger_middle_line.append(np.array(middle_pt))
        finger_lines.append(finger_middle_line)
    # cv2.imshow('intersect', intersect_result[boundy:(boundy+boundh), boundx:(boundx+boundw)])
    # cv2.waitKey()

    profile = np.zeros(img.shape[0:2], dtype=np.uint8)
    cv2.drawContours(profile, [cnts], 0, 255, 2)
    finger_width = np.zeros(5)
   
    for i in range(5):
        finger_tip = finger_poly[0, 2*i + 1, :]
        cv2.circle(intersect_result, tuple(np.int32(finger_tip)), 5, 255, 2)
        
        finger_step_width = []
        finger_middle_line = finger_lines[i]
        for j in range(len(finger_middle_line) - 2):
            finger_test[:,:] = 0        
            pos = finger_middle_line[j + 1]
            cv2.line(img, tuple(pos), tuple(finger_middle_line[j]), (255,255,255), 2)
            cv2.circle(intersect_result, tuple(np.int32(pos)), j, 255, 2)

            finger_dir = normalize(pos - finger_middle_line[j]) + normalize(finger_middle_line[j + 2] - pos)
            finger_dir = normalize(finger_dir)
            perp_dir = finger_dir.copy()
            perp_dir[0], perp_dir[1] = perp_dir[1], -perp_dir[0]
            assert np.dot(perp_dir, finger_dir) == 0

            cv2.line(finger_test, tuple(np.int32(pos + perp_dir * max_finger_half_width)),\
                tuple(np.int32(pos - perp_dir * max_finger_half_width)), 255)

            cv2.line(intersect_result, tuple(np.int32(pos + perp_dir * max_finger_half_width)),\
                tuple(np.int32(pos - perp_dir * max_finger_half_width)), 255)
            intersect = finger_test * profile
            inter_pts = intersect.nonzero()
            if inter_pts is None or inter_pts[0].shape[0] < 2:
                continue

            # cv2.imshow('intersect', intersect_result[boundy:(boundy+boundh), boundx:(boundx+boundw)])
            # cv2.waitKey()
            pt_dist = np.zeros(inter_pts[0].shape[0])
            pts_array = np.zeros((inter_pts[0].shape[0], 2), dtype=np.int32)
            for k in range(inter_pts[0].shape[0]):
                pts_array[k, :] = np.int32((inter_pts[1][k], inter_pts[0][k]))
                pt_dist[k] = np.dot(pts_array[k, :] - pos, perp_dir)

            abs_dist = np.abs(pt_dist)
            pt1_id = abs_dist.argmin()
            pt1 = pts_array[pt1_id]
            for k in range(abs_dist.shape[0]):
                pt = pts_array[k]
                if np.linalg.norm(pt - pt1) <= 4:
                    abs_dist[k] = abs_dist.max()

            pt2_id = abs_dist.argmin()
            pt2 = pts_array[pt2_id]
            if pt_dist[pt1_id] * pt_dist[pt2_id] < 0:
                width = abs(pt_dist[pt1_id] - pt_dist[pt2_id])
                width = width / coin_radius / 2 * coin_diameter_mm
                # diameter = w / (1 + 1 / (2 * d / w - 1) + (1 / (2 * d / w - 1)) ^ 2 / 2)
                d2w = args.cam_dist * 2 / width
                d2w = 1 / (d2w * 2 -1)
                width = width / (1 + d2w + d2w * d2w / 2)
                finger_step_width.append(width)
                cv2.line(img, tuple(pt1), tuple(pt2), (255, 255, 255))
                if pt_dist[pt1_id] < 0:
                    txt_pos = pt1
                else:
                    txt_pos = pt2
                cv2.putText(img, '%smm' % (round(width * 100) / 100), (txt_pos[0] + 5, txt_pos[1] + 30), font, 0.5, (255,255,255),2)
         
            cv2.circle(img, tuple(pt1), 4, (255, 255, 255), 2)
            cv2.circle(img, tuple(pt2), 4, (255, 255, 255), 2)
            

        # finger_step_width = finger_step_width[:-1]
        finger_step_width = np.array(finger_step_width)
        finger_step_width = finger_step_width[finger_step_width.shape[0] * 6 // 8:-2]
        finger_width[i] = finger_step_width.mean()
        cv2.putText(img, finger_name[i], (finger_tip[0] + 10, finger_tip[1]), font,1,(255,255,255),2)
        cv2.putText(img, '%smm' % (round(finger_width[i] * 100) / 100), (finger_tip[0] + 10, finger_tip[1] + 30), font,1,(255,255,255),2)
                
    cv2.imshow('intersect', intersect_result[boundy:(boundy+boundh), boundx:(boundx+boundw)])
    cv2.waitKey()
    


# cv2.drawContours(img,[hull],-1,(255,255,255),2)

for i in range(5):
    print('%s: %smm' % (finger_name[i], (round(finger_width[i] * 100) / 100)))
##### Show final image ########
hand_str = 'Right'
if left_hand:
    hand_str = 'Left'
cv2.putText(img, hand_str, (boundx + margin, boundy + margin),font,1,(255,255,255),2)
cv2.imshow('Result',img[boundy:(boundy+boundh), boundx:(boundx+boundw), :])
# cv2.imshow('raw', raw_img[boundy:(boundy+boundh), boundx:(boundx+boundw), :])
cv2.waitKey()
cv2.imwrite('result.jpg', img[boundy:(boundy+boundh), boundx:(boundx+boundw), :])